package com.yololabelmaker;

import javax.swing.*;
import java.awt.*;
import java.awt.event.*;
import java.awt.image.BufferedImage;
import java.io.BufferedReader;
import java.io.File;
import java.io.FileReader;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

public class CImgArea extends JPanel {
    // 根据BDD图定义属性
    private CImgOpener m_imgOpener; // properties: m_imgOpener: CImgOpener
    
    // 图片和标注相关属性
    private BufferedImage currentImage;
    private String currentImagePath;
    private List<LabelRect> labelRects;
    private LabelRect currentRect;
    private Point startPoint;
    private Point endPoint;
    private boolean isDrawing;
    private boolean fitToWindow;
    private double scaleFactor;
    
    public CImgArea() {
        m_imgOpener = new CImgOpener();
        labelRects = new ArrayList<>();
        isDrawing = false;
        fitToWindow = true;
        scaleFactor = 1.0;
        
        // 初始化拖拽支持
        setTransferHandler(new ImageTransferHandler(this));
        
        // 添加鼠标监听器
        addMouseListener(new MouseAdapter() {
            @Override
            public void mousePressed(MouseEvent e) {
                if (currentImage != null) {
                    startPoint = adjustForScale(e.getPoint());
                    currentRect = new LabelRect();
                    isDrawing = true;
                }
            }
            
            @Override
            public void mouseReleased(MouseEvent e) {
                if (isDrawing && currentImage != null) {
                    endPoint = adjustForScale(e.getPoint());
                    // 确保矩形有效（宽高大于10像素）
                    if (Math.abs(endPoint.x - startPoint.x) > 10 && Math.abs(endPoint.y - startPoint.y) > 10) {
                        currentRect.setRect(startPoint, endPoint);
                        
                        // 显示对话框让用户输入类别
                        String className = JOptionPane.showInputDialog("请输入类别名称：");
                        if (className != null && !className.trim().isEmpty()) {
                            currentRect.setClassName(className);
                            labelRects.add(currentRect);
                        }
                    }
                    isDrawing = false;
                    repaint();
                }
            }
        });
        
        // 添加鼠标移动监听器
        addMouseMotionListener(new MouseMotionAdapter() {
            @Override
            public void mouseDragged(MouseEvent e) {
                if (isDrawing && currentImage != null) {
                    endPoint = adjustForScale(e.getPoint());
                    repaint();
                }
            }
        });
        
        // 添加组件监听器以处理窗口大小改变
        addComponentListener(new ComponentAdapter() {
            @Override
            public void componentResized(ComponentEvent e) {
                if (fitToWindow && currentImage != null) {
                    calculateScaleFactor();
                    repaint();
                }
            }
        });
        
        // 设置背景色
        setBackground(Color.LIGHT_GRAY);
    }
    
    @Override
    protected void paintComponent(Graphics g) {
        super.paintComponent(g);
        
        if (currentImage != null) {
            // 绘制图片
            int width = (int) (currentImage.getWidth() * scaleFactor);
            int height = (int) (currentImage.getHeight() * scaleFactor);
            g.drawImage(currentImage, 0, 0, width, height, null);
            
            // 绘制已存在的标签框
            Graphics2D g2d = (Graphics2D) g;
            g2d.setColor(Color.RED);
            g2d.setStroke(new BasicStroke(2));
            
            for (LabelRect rect : labelRects) {
                int x = (int) (rect.getX() * scaleFactor);
                int y = (int) (rect.getY() * scaleFactor);
                int w = (int) (rect.getWidth() * scaleFactor);
                int h = (int) (rect.getHeight() * scaleFactor);
                
                g2d.drawRect(x, y, w, h);
                
                // 绘制类别名称
                g2d.setColor(Color.YELLOW);
                g2d.fillRect(x, y - 20, w, 20);
                g2d.setColor(Color.BLACK);
                g2d.drawString(rect.getClassName(), x + 5, y - 5);
            }
            
            // 绘制正在绘制的标签框
            if (isDrawing && startPoint != null && endPoint != null) {
                g2d.setColor(Color.GREEN);
                g2d.setStroke(new BasicStroke(2));
                
                int x = Math.min((int)(startPoint.x * scaleFactor), (int)(endPoint.x * scaleFactor));
                int y = Math.min((int)(startPoint.y * scaleFactor), (int)(endPoint.y * scaleFactor));
                int w = Math.abs((int)((endPoint.x - startPoint.x) * scaleFactor));
                int h = Math.abs((int)((endPoint.y - startPoint.y) * scaleFactor));
                
                g2d.drawRect(x, y, w, h);
            }
        }
    }
    
    // 设置图片
    public void setImage(BufferedImage image, String path) {
        this.currentImage = image;
        this.currentImagePath = path;
        this.labelRects.clear();
        
        // 如果设置了适应窗口，则计算缩放因子
        if (fitToWindow) {
            calculateScaleFactor();
        } else {
            scaleFactor = 1.0;
        }
        
        // 尝试加载已存在的标签文件
        loadExistingLabels();
        
        repaint();
    }
    
    // 清除标签
    public void clearLabels() {
        labelRects.clear();
        repaint();
    }
    
    // 保存标签
    public void saveLabels() {
        if (currentImage != null && currentImagePath != null && !labelRects.isEmpty()) {
            YOLOLabelSaver saver = new YOLOLabelSaver();
            saver.saveLabels(currentImagePath, currentImage, labelRects);
        }
    }
    
    // 加载已存在的标签文件
    private void loadExistingLabels() {
        if (currentImagePath != null && currentImage != null) {
            YOLOLabelLoader loader = new YOLOLabelLoader();
            List<LabelRect> loadedRects = loader.loadLabels(currentImagePath);
            if (loadedRects != null && !loadedRects.isEmpty()) {
                // 尝试从标签文件中读取完整的YOLO格式数据
                File imageFile = new File(currentImagePath);
                String parentDir = imageFile.getParent();
                String imageFileName = imageFile.getName();
                String labelFileName = imageFileName.substring(0, imageFileName.lastIndexOf('.')) + ".txt";
                String labelFilePath = "";
                
                // 尝试不同位置查找标签文件
                if (parentDir != null && parentDir.endsWith("images")) {
                    labelFilePath = parentDir.replace("images", "labels") + File.separator + labelFileName;
                } else if (parentDir != null) {
                    labelFilePath = parentDir + File.separator + "labels" + File.separator + labelFileName;
                    if (!new File(labelFilePath).exists()) {
                        labelFilePath = parentDir + File.separator + labelFileName;
                    }
                }
                
                // 如果是网络图片，从当前目录的labels文件夹加载
                if (currentImagePath.startsWith("network_image_")) {
                    labelFilePath = System.getProperty("user.dir") + File.separator + "labels" + File.separator + labelFileName;
                }
                
                File labelFile = new File(labelFilePath);
                if (labelFile.exists()) {
                    try (BufferedReader reader = new BufferedReader(new FileReader(labelFile))) {
                        String line;
                        int index = 0;
                        while ((line = reader.readLine()) != null && index < loadedRects.size()) {
                            line = line.trim();
                            if (line.isEmpty() || line.startsWith("#")) {
                                continue;
                            }
                            
                            // 解析YOLO格式的坐标
                            String[] parts = line.split("\\s+");
                            if (parts.length >= 5) {
                                try {
                                    double centerX = Double.parseDouble(parts[1]);
                                    double centerY = Double.parseDouble(parts[2]);
                                    double width = Double.parseDouble(parts[3]);
                                    double height = Double.parseDouble(parts[4]);
                                    
                                    // 更新标签框的坐标
                                    LabelRect rect = loadedRects.get(index);
                                    YOLOLabelLoader.updateLabelRectCoordinates(
                                            rect, currentImage.getWidth(), currentImage.getHeight(), 
                                            centerX, centerY, width, height);
                                } catch (NumberFormatException e) {
                                    // 忽略解析错误的行
                                }
                            }
                            index++;
                        }
                    } catch (IOException e) {
                        // 忽略读取错误
                    }
                }
                
                labelRects.addAll(loadedRects);
            }
        }
    }
    
    // 计算缩放因子
    private void calculateScaleFactor() {
        if (currentImage != null) {
            double scaleX = (double) getWidth() / currentImage.getWidth();
            double scaleY = (double) getHeight() / currentImage.getHeight();
            scaleFactor = Math.min(scaleX, scaleY);
        }
    }
    
    // 根据缩放调整坐标
    private Point adjustForScale(Point point) {
        int x = (int) (point.x / scaleFactor);
        int y = (int) (point.y / scaleFactor);
        return new Point(x, y);
    }
    
    // 检查是否有图片加载
    public boolean hasImage() {
        return currentImage != null;
    }
    
    // 设置是否适应窗口
    public void setFitToWindow(boolean fitToWindow) {
        this.fitToWindow = fitToWindow;
        if (currentImage != null) {
            if (fitToWindow) {
                calculateScaleFactor();
            } else {
                scaleFactor = 1.0;
            }
            repaint();
        }
    }
    
    // 获取当前缩放因子
    public double getScaleFactor() {
        return scaleFactor;
    }
    
    // 设置缩放因子
    public void setScaleFactor(double scaleFactor) {
        this.scaleFactor = scaleFactor;
        this.fitToWindow = false;
        repaint();
    }
}